#include "CNN/cifar10.hpp"
#include "CNN/cifar10_model.hpp"

const std::string kRoot = "build/data/cifar-10-batches-bin";

// x: NCHW
// w: OIHW  (x:C == w:I)
torch::Tensor ele_conv(const torch::Tensor &x, const torch::Tensor &w,
                       const int stride) {
  using namespace torch::indexing;
  assert(x.size(2) == x.size(3) && "CIFAR-10:images : H = W = 32");
  assert(w.size(2) == w.size(3) && "CIFAR-10:filter : K * K");
  // L: input length
  const auto N = x.size(0), I = x.size(1), L = x.size(2);
  const auto O = w.size(0), K = w.size(2);

  assert(((L - K) % stride) == 0);
  const auto lo = (L - K) / stride + 1; // output length
  const auto Fin = lo * lo * O;

  // - FC.weight: Fout * Fin
  // - FC: X @ transpose(FC.weight), where X \in R^{N * Fin}
  // - X <=> conv_out
  torch::Tensor out = torch::zeros({N, Fin});
  for (int in = 0; in < I; in++) {
    for (int i = 0; i < K; i++) {
      for (int j = 0; j < K; j++) {
        // Get input:
        //  1. (N, lo, lo)
        torch::Tensor xx = x.as_strided(
            {N, lo, lo},
            {x.stride(0), x.stride(2) * stride, x.stride(3) * stride},
            in * L * L + i * x.stride(2) + j);
        //  2. (N, Fin)
        xx = xx.reshape({N, 1, lo * lo})
                 .broadcast_to({N, O, lo * lo})
                 .reshape({N, Fin});

        // Construct filter: (N, Fin)
        torch::Tensor f = w.index({Slice(), in, i, j})
                              .reshape({O, 1})
                              .broadcast_to({O, lo * lo})
                              .flatten()
                              .reshape({1, Fin})
                              .broadcast_to({N, Fin});
        out += xx * f;
      }
    }
  }

  return out.reshape({N, O, lo, lo});
}

void toy_example_check(int n) {
  torch::manual_seed(0);
  NaiveNoPaddingNet5 model;

  const std::string pt_file = kRoot + "/" + model.pt_name;
  torch::serialize::InputArchive archive;
  archive.load_from(pt_file);
  model.load(archive);

  // * err: 3.68649e-06
  torch::Tensor x = torch::rand({n, 3, 32, 32});

  // clang-format off
  // ! err: 0.0090315 ~ 1e-2
  // torch::Tensor x = torch::range(1, n * 3 * 32 * 32, 1).reshape({n, 3, 32, 32});
  // clang-format on

  auto expected = model.conv(x);
  auto res =
      ele_conv(x, model.conv->weight, NaiveNoPaddingNet5::CIFAR10_STRIDE);

  std::cout << torch::norm(expected - res) << std::endl;
}

int main() {
  toy_example_check(1);
  toy_example_check(2);
  return 0;
}